import torch.nn as nn

from mmdet.models import NECKS

from models.backbones.shufflenetv2 import InvertedResidual


@NECKS.register_module()
class SSDLiteNeckShuffleNetv2(nn.Module):
    def __init__(self, in_channel: int, extra_channels: list):
        super().__init__()
        assert len(extra_channels) == 3
        self.extras = nn.ModuleList([
            InvertedResidual(inp=in_channel,
                             oup=extra_channels[0],
                             stride=2,
                             padding=1),
            InvertedResidual(inp=extra_channels[0],
                             oup=extra_channels[1],
                             stride=2,
                             padding=1),
            InvertedResidual(inp=extra_channels[1],
                             oup=extra_channels[2],
                             stride=2,
                             padding=0)
        ])

    def forward(self, x):
        outs = list()
        for extra in self.extras:
            x = extra(x)
            outs.append(x)
        return outs
